package com.lyh.notes.controller;


import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.lyh.notes.model.base.ApiResponse;
import com.lyh.notes.model.vo.ai.AiVO;
import com.lyh.notes.service.OpenAiService;
import io.reactivex.Flowable;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.validation.Valid;
import java.io.IOException;
import java.util.Arrays;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;


@RestController
public class OpenAiController {

    @Autowired
    private OpenAiService openAiService;

    @PostMapping("/api/ai")
    public ApiResponse<String> askQuestion(@RequestBody@Valid AiVO aiVO) {
        return openAiService.generateText(aiVO.getPrompt());
    }

    @PostMapping("/api/ai-flux")
    public SseEmitter streamCompletion(@RequestParam String prompt) {
        SseEmitter emitter = new SseEmitter(30_000L); // 30秒超时

        // 异步处理流式响应
        CompletableFuture.runAsync(() -> {
            try {
                Generation gen = new Generation();
                Message userMsg = Message.builder()
                        .role(Role.USER.getValue())
                        .content(prompt)
                        .build();

                GenerationParam param = GenerationParam.builder()
                        .apiKey(System.getenv("sk-3cc34942bc9f494787a2f3daf7623d21"))
                        .model("qwen-plus")
                        .messages(Arrays.asList(userMsg))
                        .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                        .incrementalOutput(true)
                        .build();

                Flowable<GenerationResult> result = gen.streamCall(param);
                result.blockingForEach(message -> {
                    String content = message.getOutput().getChoices().get(0).getMessage().getContent();
                    try {
                        emitter.send(SseEmitter.event()
                                .data(content)
                                .id(UUID.randomUUID().toString()));
                    } catch (IOException e) {
                        throw new RuntimeException("SSE发送失败", e);
                    }
                });

                emitter.complete();
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
        });

        return emitter;
    }
}
